{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "license"
      },
      "source": [
        "##### *Copyright 2020 Google LLC*\n",
        "*Licensed under the Apache License, Version 2.0 (the \"License\")*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "rKwqeqWBXANA"
      },
      "outputs": [],
      "source": [
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hRTa3Ee15WsJ"
      },
      "source": [
        "# Fix top issues when converting a TF model for Edge TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TaX0smDP7xQY"
      },
      "source": [
        "This page shows how to fix some known issues when converting TensorFlow 2 models for the Edge TPU. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "viewin-badges"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/google-coral/tutorials/blob/master/fix_conversion_issues_ptq_tf2.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\"\u003e\u003c/a\u003e\n",
        "\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\n",
        "\u003ca href=\"https://github.com/google-coral/tutorials/blob/master/fix_conversion_issues_ptq_tf2.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://img.shields.io/static/v1?logo=GitHub\u0026label=\u0026color=333333\u0026style=flat\u0026message=View%20on%20GitHub\" alt=\"View in GitHub\"\u003e\u003c/a\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "BnSreNhbCQ69"
      },
      "source": [
        "To run all the code in this tutorial, select **Runtime \u003e Run all** in the Colab toolbar."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GTCYQg_be8C0"
      },
      "source": [
        "## Set up the environment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "02MxhCyFmpzn"
      },
      "source": [
        "Import the Python libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "iBMcobPHdD8O"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "assert float(tf.__version__[:3]) \u003e= 2.3\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0E5ZWVu7BiCk"
      },
      "source": [
        "Install the Edge TPU Compiler:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "p6ZpWgrk21Ad"
      },
      "outputs": [],
      "source": [
        "! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n",
        "\n",
        "! echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
        "\n",
        "! sudo apt-get update\n",
        "\n",
        "! sudo apt-get install edgetpu-compiler\t"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SFabFnqcH3jc"
      },
      "source": [
        "### Create quantization function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6MPyPU7RH7lO"
      },
      "outputs": [],
      "source": [
        "def quantize_model(converter):\n",
        "  # This generator provides a junk representative dataset\n",
        "  # (It creates a poor model but is only for demo purposes)\n",
        "  def representative_data_gen():\n",
        "    for i in range(10):\n",
        "      image = tf.random.uniform([1, 224, 224, 3])\n",
        "      yield [image]\n",
        "    \n",
        "  converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "  converter.representative_dataset = representative_data_gen\n",
        "  converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "  converter.inference_input_type = tf.uint8\n",
        "  converter.inference_output_type = tf.uint8\n",
        "\n",
        "  return converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QpvppqCHdXvc"
      },
      "source": [
        "## Can't compile due to dynamic batch size"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qiBIONcBgC44"
      },
      "source": [
        "The Edge TPU Compiler fails for some models such as MobileNetV1 if the input shape batch size is not set to 1, although this isn't exactly obvious from the compiler's output:\n",
        "\n",
        "```\n",
        "Invalid model: mobilenet_quant.tflite\n",
        "Model not quantized\n",
        "```\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "L9NR3xv3Tr0k"
      },
      "source": [
        "That error might be caused by something else, but you should try the following solution because although it's not required for all models, it shouldn't hurt."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jMenehKsXoLf"
      },
      "source": [
        "### Solution for a Keras model object"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "2FS5RMZ1Fmg6"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.applications.MobileNet()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hNvMl6CM6lG4"
      },
      "source": [
        "The following creates a TFLite file that will fail in the Edge TPU Compiler:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "srOYhMYfx9XH"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "tflite_model = quantize_model(converter)\n",
        "\n",
        "with open('mobilenet_quant_before.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "joxrIB0I3cdi"
      },
      "outputs": [],
      "source": [
        "! edgetpu_compiler mobilenet_quant_before.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QqiCWFJDQo4q"
      },
      "source": [
        "It won't compile because the model has a dynamic batch size as shown here (`None` means it's dynamic):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yyP5h8UNQuRE"
      },
      "outputs": [],
      "source": [
        "model.input.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7qtCQN0QQ2QR"
      },
      "source": [
        "So to fix it, we need to set that to 1:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1_ViBvedQ6SK"
      },
      "outputs": [],
      "source": [
        "model.input.set_shape((1,) + model.input.shape[1:])\n",
        "model.input.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mtPcYiER3Ymp"
      },
      "source": [
        "Now we can convert it again and it will compile:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_eal9xLNQ9pH"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "tflite_model = quantize_model(converter)\n",
        "\n",
        "with open('mobilenet_quant_after.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "eF9PeeNOL-y5"
      },
      "outputs": [],
      "source": [
        "! edgetpu_compiler mobilenet_quant_after.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4vP9QqNeXd37"
      },
      "source": [
        "### Solution for a SavedModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vMeIbwBncNQE"
      },
      "source": [
        "If you're loading a SavedModel file, then the fix looks a little different.\n",
        "\n",
        "So let's say you saved a model like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "G1kCymYOYHAc"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.applications.MobileNet()\n",
        "save_path = os.path.join(\"mobilenet/1/\")\n",
        "tf.saved_model.save(model, save_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "odQFGxlMatPu"
      },
      "source": [
        "Ideally, you could later load the model like this:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "KqArhmo5WwiW"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_saved_model(save_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_q7AYuPvW0Sd"
      },
      "source": [
        "But the saved model's input still has a dynamic batch size, so you need to instead load the model with `saved_model.load()` and modify the input's concrete function so it has a batch size of 1. Then load it into `TFLiteConverter` using the concrete function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kujXSsDoXxCr"
      },
      "outputs": [],
      "source": [
        "imported = tf.saved_model.load(save_path)\n",
        "concrete_func = imported.signatures[\"serving_default\"]\n",
        "concrete_func.inputs[0].set_shape([1, 224, 224, 3])\n",
        "converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "p1MX3eEOX27S"
      },
      "source": [
        "Now you can convert to TFLite and it will compile for the Edge TPU:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vExrD69bWTRc"
      },
      "outputs": [],
      "source": [
        "tflite_model = quantize_model(converter)\n",
        "with open('mobilenet_imported_quant.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sbYbcGJhWmnm"
      },
      "outputs": [],
      "source": [
        "! edgetpu_compiler mobilenet_imported_quant.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "N5xsHyRDCbCv"
      },
      "source": [
        "## Can't import a SavedModel without a signature"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jyYoNMGtcVJV"
      },
      "source": [
        "Sometimes, a SavedModel does not include a signature (such as when the model was built with a custom `tf.Module`), making it impossible to load using `TFLiteConverter`. So you can instead add the batch size as follows.\n",
        "\n",
        "**Note:** If you created the model yourself, see how to [specify the signature during export](https://www.tensorflow.org/guide/saved_model#specifying_signatures_during_export) so this isn't a problem."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8Jwb_0xjcv1G"
      },
      "outputs": [],
      "source": [
        "# First get the Inception SavedModel, which is lacking a signature\n",
        "!wget -O imagenet_inception_v2_classification_4.tar.gz https://tfhub.dev/google/imagenet/inception_v2/classification/4?tf-hub-format=compressed\n",
        "!mkdir -p imagenet_inception_v2_classification_4\n",
        "!tar -xvzf imagenet_inception_v2_classification_4.tar.gz --directory imagenet_inception_v2_classification_4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Gs6IpZgmDGZx"
      },
      "source": [
        "For example, this fails because the model has no signature:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JOl64eEzZpSH"
      },
      "outputs": [],
      "source": [
        "#converter = tf.lite.TFLiteConverter.from_saved_model(\"imagenet_inception_v2_classification_4\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bIyHbP2WDJDj"
      },
      "source": [
        "Whereas other code above loads the input's concrete function by calling upon its \"serving_default\" signature, we can't do that if the model has no signature. So we instead get the concrete function by specifying its known input tensor shape:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NK05ZlMocy0E"
      },
      "outputs": [],
      "source": [
        "imported = tf.saved_model.load(\"imagenet_inception_v2_classification_4\")\n",
        "concrete_func = imported.__call__.get_concrete_function(\n",
        "          tf.TensorSpec([1, 224, 224, 3]))\n",
        "converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7MSnuIdeYak7"
      },
      "source": [
        "Now we can convert and compile:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dPrn8OfyYbHh"
      },
      "outputs": [],
      "source": [
        "tflite_model = quantize_model(converter)\n",
        "with open('inceptionv2_imported_quant.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t8Daq2zsZ5yX"
      },
      "outputs": [],
      "source": [
        "! edgetpu_compiler inceptionv2_imported_quant.tflite"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "license"
      ],
      "name": "fix_conversion_issues_ptq_tf2.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
